refactor: Refactor MockTexeraDB into a JVM singleton#6238
Conversation
Each spec that mixes in MockTexeraDB used to start its own EmbeddedPostgres in beforeAll. Once the global Tags.limit(Tags.Test, 1) restriction was lifted, these specs ran in parallel and collided on the shared /tmp/embedded-pg/<binary>/ directory with an OverlappingFileLockException at EmbeddedPostgres.prepareBinaries. Move the EmbeddedPostgres + DSLContext fields onto a companion object, make ensureInitialized() idempotent and synchronized, and turn shutdownDB() into a no-op so the singleton lives for the JVM. The trait still exposes the same methods so spec callers do not change.
Automated Reviewer SuggestionsBased on the
|
Neilk1021
changed the title
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #6238 +/- ##
============================================
+ Coverage 56.89% 59.77% +2.88%
- Complexity 3058 3370 +312
============================================
Files 1129 1128 -1
Lines 43802 43856 +54
Branches 4743 4753 +10
============================================
+ Hits 24922 26217 +1295
+ Misses 17449 16207 -1242
- Partials 1431 1432 +1
*This pull request uses carry forward flags. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
| config | throughput | MB/s | latency | max Δ latest / 7d | |
|---|---|---|---|---|---|
| 🔴 | bs=10 sw=10 sl=64 | 394 | 0.24 | 24,622/37,341/37,341 us | 🔴 +19.4% / 🔴 +138.6% |
| 🔴 | bs=100 sw=10 sl=64 | 798 | 0.487 | 124,993/135,105/135,105 us | 🟢 -12.4% / 🔴 +23.8% |
| ⚪ | bs=1000 sw=10 sl=64 | 897 | 0.548 | 1,105,152/1,219,585/1,219,585 us | ⚪ within ±5% / 🔴 +16.6% |
Baseline details
Latest main 7c135ee from same runner
| config | metric | PR | latest main | 7d avg | Δ latest | Δ 7d |
|---|---|---|---|---|---|---|
| bs=10 sw=10 sl=64 | throughput | 394 tuples/sec | 416 tuples/sec | 760.25 tuples/sec | -5.3% | -48.2% |
| bs=10 sw=10 sl=64 | MB/s | 0.24 MB/s | 0.254 MB/s | 0.464 MB/s | -5.5% | -48.3% |
| bs=10 sw=10 sl=64 | p50 | 24,622 us | 23,267 us | 12,873 us | +5.8% | +91.3% |
| bs=10 sw=10 sl=64 | p95 | 37,341 us | 31,271 us | 15,648 us | +19.4% | +138.6% |
| bs=10 sw=10 sl=64 | p99 | 37,341 us | 31,271 us | 19,338 us | +19.4% | +93.1% |
| bs=100 sw=10 sl=64 | throughput | 798 tuples/sec | 828 tuples/sec | 974.75 tuples/sec | -3.6% | -18.1% |
| bs=100 sw=10 sl=64 | MB/s | 0.487 MB/s | 0.506 MB/s | 0.595 MB/s | -3.8% | -18.1% |
| bs=100 sw=10 sl=64 | p50 | 124,993 us | 118,980 us | 102,470 us | +5.1% | +22.0% |
| bs=100 sw=10 sl=64 | p95 | 135,105 us | 154,259 us | 109,144 us | -12.4% | +23.8% |
| bs=100 sw=10 sl=64 | p99 | 135,105 us | 154,259 us | 115,530 us | -12.4% | +16.9% |
| bs=1000 sw=10 sl=64 | throughput | 897 tuples/sec | 910 tuples/sec | 1,000 tuples/sec | -1.4% | -10.3% |
| bs=1000 sw=10 sl=64 | MB/s | 0.548 MB/s | 0.555 MB/s | 0.61 MB/s | -1.3% | -10.2% |
| bs=1000 sw=10 sl=64 | p50 | 1,105,152 us | 1,097,711 us | 1,005,545 us | +0.7% | +9.9% |
| bs=1000 sw=10 sl=64 | p95 | 1,219,585 us | 1,205,873 us | 1,045,968 us | +1.1% | +16.6% |
| bs=1000 sw=10 sl=64 | p99 | 1,219,585 us | 1,205,873 us | 1,076,800 us | +1.1% | +13.3% |
Raw CSV
config_idx,batch_size,schema_width,string_len,num_batches,total_ms,total_tuples,total_bytes,tuples_per_sec,mb_per_sec,lat_p50_us,lat_p95_us,lat_p99_us
0,10,10,64,20,507.89,200,128000,394,0.240,24622.31,37341.21,37341.21
1,100,10,64,20,2506.78,2000,1280000,798,0.487,124993.34,135104.50,135104.50
2,1000,10,64,20,22289.98,20000,12800000,897,0.548,1105151.89,1219585.40,1219585.40
|
I noticed that the withTransaction block in DatasetResource.scala has gotten a bit bloated/heavy, which is likely contributing to the latency spikes seen in the benchmarks. I'm planning to submit a follow-up fix to slim down the transaction scope, but wanted to call it out here. |
You can open an issue for this. |
|
@Yicong-Huang, can you take a look at Benchmark? Is it supposed to be better? |
There was a problem hiding this comment.
Pull request overview
Refactors the test database infrastructure to use a JVM-wide singleton EmbeddedPostgres, provisioning per-suite isolated databases to reduce test runtime and avoid concurrent extraction/file-lock issues, while also adjusting transaction handling in SqlServer.withTransaction to better tolerate mixed autoCommit states seen in concurrent tests.
Changes:
- Converted
MockTexeraDBinto a singleton-backed design with cached DDL loading and per-suite database creation. - Updated multipart-upload concurrency tests to use independent JDBC connections against the per-suite database.
- Reworked
SqlServer.withTransactionto manually manageautoCommit/ commit / rollback behavior.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| file-service/src/test/scala/org/apache/texera/service/resource/DatasetResourceSpec.scala | Switches lock-simulation tests to use independent JDBC connections from the singleton embedded Postgres. |
| common/dao/src/test/scala/org/apache/texera/dao/MockTexeraDB.scala | Introduces JVM-singleton EmbeddedPostgres, caches cleaned DDL, creates per-suite databases, and manages per-test cleanup. |
| common/dao/src/main/scala/org/apache/texera/dao/SqlServer.scala | Replaces jOOQ-managed transactions with manual connection/autoCommit handling in withTransaction. |
| amber/src/test/scala/org/apache/texera/amber/engine/e2e/TestUtils.scala | Updates e2e DB initialization to use the singleton embedded Postgres + cached DDL. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| } finally { | ||
| try { | ||
| if (!conn.isClosed) { | ||
| scala.util.Using.resource(conn.createStatement()) { stmt => | ||
| stmt.execute( | ||
| """ | ||
| DO $$ DECLARE | ||
| r RECORD; | ||
| BEGIN | ||
| FOR r IN (SELECT tablename FROM pg_tables WHERE schemaname = 'public') LOOP | ||
| EXECUTE 'TRUNCATE TABLE ' || quote_ident(r.tablename) || ' CASCADE'; | ||
| END LOOP; | ||
| END $$; | ||
| """ | ||
| ) | ||
| } | ||
| } | ||
| } catch { | ||
| case e: Exception => e.printStackTrace() | ||
| } | ||
| } |
| def getDDLScript: String = ddlScript.getOrElse(throw new RuntimeException("DDL not loaded")) | ||
| } | ||
|
|
||
| trait MockTexeraDB extends AnyFlatSpecLike { |
| val txCtx = org.jooq.impl.DSL.using(conn, dsl.dialect()) | ||
| val result = block(txCtx) | ||
|
|
||
| if (originalAutoCommit) { | ||
| conn.commit() | ||
| } | ||
|
|
||
| result | ||
| } catch { | ||
| case e: Throwable => | ||
| if (originalAutoCommit) { | ||
| conn.rollback() | ||
| } | ||
| throw e | ||
| } finally { | ||
| try { | ||
| if (originalAutoCommit != conn.getAutoCommit) { | ||
| conn.setAutoCommit(originalAutoCommit) | ||
| } | ||
| } finally { | ||
| provider.release(conn) | ||
| } | ||
| } |
please ignore benchmark. it is quite not stable. until we stabilize it, the signals are informational only. |
|
@mengw15 Please review it as well. |
Adds proper guards to MockTexeraDB initialization to guard against the DDL script not being loaded. Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Neil Ketteringham <53205839+Neilk1021@users.noreply.github.com>
…se where a test disables autocommit.
…exeraDB's single connection with Hikari Connection pooling
|
Thanks @Neilk1021 for the PR, please wait for my review, might need 1 day or 2 to find time. meanwhile, I hope to loop in @mengw15 as he is modifying CI as well for another effort. @mengw15 please post your related issues here to also give @Neilk1021 some context. |
|
Hi @Neilk1021 👋 — for context, here's the CI work I have in flight (in
Looks like we're on different layers — you're optimizing the embedded unit-test DB ( |
|
For future reviewers, small notes on why I reverted my changes to with transaction and implemented Hikari pooling for First, my earlier The real root cause: backing On test speed: the physical connection was never the bottleneck. The expensive part was spinning up |
Thanks for the context! |
|
@mengw15 @Neilk1021 : as you two work very "closely," feel free to have offline discussions and report the results on this PR. |
Xiao-zhen-Liu
left a comment
Xiao-zhen-Liu
left a comment
There was a problem hiding this comment.
Thanks for taking this on — the core design is a clear improvement: one embedded Postgres per JVM, a fresh database per suite, cached DDL, and a real pool are the right shape, and the startup-time win is real. I ran the changed e2e specs on JDK 11 to confirm the embedded-DB switch is safe (note on TestUtils below), and it holds up.
One thing I think should block merge: the per-test cleanup in withFixture doesn't actually run. It truncates the public schema, but all our tables are in the texera_db schema (sql/texera_ddl.sql:42-43), so it matches zero tables. I confirmed by running ComputingUnitAccessSpec and FileResolverSpec on this branch — both pass, but only because the beforeAll seed is never deleted; when I pointed the truncation at the real schema, ComputingUnitAccessSpec failed 3/5. So "each test starts clean" isn't true today, and it can't be fixed by just changing the schema name without also moving several suites' beforeAll seeding into beforeEach. I'd suggest either dropping the per-test truncation and keeping each suite's own cleanup, or making it real and fixing the affected suites here.
Two notes so the description stays accurate: it mentions a SqlServer.withTransaction change, but SqlServer.scala isn't in the diff — the autoCommit fix actually comes from the new pool plus the DatasetResourceSpec connection change. And this is good groundwork toward parallel suites (#4525) but doesn't get us there yet: production code still reads through the single global SqlServer, which only works while Tags.limit keeps tests sequential. Worth saying that explicitly, and that the e2e embedded-DB switch depends on #4179 having landed. Smaller cleanups are inline.
What's still TODO for the series: (1) decide + fix the per-test cleanup; (2) make the SqlServer context per-suite/thread (the real prerequisite for parallel suites); (3) only then remove Tags.limit + enable parallelExecution; (4) drop created DBs, remove the now-dead jdbcUrlForTestCases config, dedupe the provisioning helper, avoid the double pool.
| DO $$ DECLARE | ||
| r RECORD; | ||
| BEGIN | ||
| FOR r IN (SELECT tablename FROM pg_tables WHERE schemaname = 'public') LOOP |
There was a problem hiding this comment.
This per-test cleanup is a silent no-op (targets the wrong schema). It truncates the public schema, but every table lives in texera_db (sql/texera_ddl.sql:42-43 creates that schema and sets search_path), so pg_tables WHERE schemaname='public' matches nothing. Verified on JDK 11: ComputingUnitAccessSpec and FileResolverSpec pass only because the beforeAll seed is never deleted; retargeting the truncation at texera_db makes ComputingUnitAccessSpec fail 3/5. Two consequences: (1) "each test starts clean" is false today — the trait adds no per-test isolation, which is a trap for anyone who later drops a suite's own cleanup trusting this; (2) you can't just fix the schema name, because that breaks every suite that seeds only in beforeAll (AccessControlResourceSpec, ComputingUnitAccessSpec, WorkflowResourceSpec, FileResolverSpec, and the DatasetResourceSpec this PR touches). Suggest either dropping the truncation and keeping each suite's existing cleanup, or making it real and moving those seeds to beforeEach.
There was a problem hiding this comment.
I think removing the truncation for now and leaving the truncation logic for later would be best? This allows us to get the speed-up now without ballooning the responsibilities of this PR?
| } | ||
| } | ||
| } catch { | ||
| case e: Exception => e.printStackTrace() |
There was a problem hiding this comment.
If the truncation throws it's only printed and swallowed, so the next test can start against dirty state. If the cleanup stays, a failure to clean should fail loudly rather than print-and-continue.
| testScopedContext = Some(scopedCtx) | ||
|
|
||
| // Point the Texera backend exactly to this suite's isolated database | ||
| SqlServer.initConnection(jdbcUrl, "postgres", "") |
There was a problem hiding this comment.
SqlServer.initConnection constructs a SqlServer whose constructor eagerly opens its own 10-connection Hikari pool (SqlServer.scala:62, minimumIdle=2), and the next line swaps in the trait's separate pool. So each trait suite opens two pools and the SqlServer-owned one is never used. Consider a SqlServer method that repoints its existing pool, or reuse that pool instead of building a new DataSourceConnectionProvider.
|
|
||
| sqlServerInstance.replaceDSLContext(dslContext.get) | ||
| try { | ||
| sqlServerInstance.replaceDSLContext(activeContext) |
There was a problem hiding this comment.
Per-suite isolation is real for getDSLContext, but production DAOs read through the single global SqlServer, which this repoints per test. That's safe only while Tags.limit(Tags.Test,1) keeps suites sequential (still in all 12 build.sbt). Once #4525 removes that, two suites in one JVM race on this global — and initConnection closing the prior instance's pool can drop a pool another suite is using. Worth noting this PR is groundwork; parallel-safety needs a follow-up that makes the context per-suite/thread.
| def initializeDBAndReplaceDSLContext(): Unit = { | ||
| assert(dbInstance.isEmpty && dslContext.isEmpty) | ||
| val parts: Array[String] = content.split("(?m)^CREATE DATABASE :\"DB_NAME\";") | ||
| val sqlBody = if (parts.length > 1) parts(1) else content |
There was a problem hiding this comment.
If the CREATE DATABASE split marker ever changes, parts.length > 1 is false and this silently runs the whole DDL as the body. A hard failure would be safer than quietly doing the wrong thing.
| sql.linesIterator.filterNot(_.trim.startsWith("\\c")).mkString("\n") | ||
|
|
||
| val embedded = EmbeddedPostgres.builder().start() | ||
| var tablesAndIndexCreation = removeCCommands(sqlBody) |
There was a problem hiding this comment.
tablesAndIndexCreation is never reassigned — can be a val.
| import scala.util.Using | ||
|
|
||
| /** | ||
| * Provides a JVM-singleton EmbeddedPostgres for tests. Multiple specs that mix |
There was a problem hiding this comment.
This credits the singleton with avoiding OverlappingFileLockException "in parallel," but tests still run sequentially (Tags.limit is in place), so the actual current win is booting Postgres once. Worth aligning the comment with today's behavior.
| testScopedContext.get | ||
| } | ||
|
|
||
| def getDBInstance: EmbeddedPostgres = MockTexeraDB.getDBInstance |
There was a problem hiding this comment.
Minor readability: getDBInstance exists on both the object and this trait, and there are two init entry points (ensureInitialized vs initializeDBAndReplaceDSLContext). A line of doc on each boundary would save a maintainer some head-scratching.
| } | ||
| } | ||
|
|
||
| val targetDbConn = embedded.getDatabase("postgres", dbName).getConnection |
There was a problem hiding this comment.
A few things in this method: (1) it duplicates the "create DB + run DDL" sequence from the trait (MockTexeraDB.scala:101-114) and the copies have already drifted — worth one shared helper on object MockTexeraDB. (2) targetDbConn is closed manually after the Using.resource, so it leaks if execute throws — wrap it in Using.resource too. (3) The DROP DATABASE IF EXISTS on a freshly generated UUID name (line 208) is dead. (4) This quietly moves the e2e specs off the external test Postgres onto the embedded one — a real (good) change worth a line in the description; it depends on #4179. I ran the e2e specs this way on JDK 11 (46 tests ×3, no Connection refused), so the switch looks safe for sequential runs.
|
|
||
| val connectionProvider = getDSLContext.configuration().connectionProvider() | ||
| val connection = connectionProvider.acquire() | ||
| val connection = MockTexeraDB.getDBInstance.getDatabase("postgres", uniqueDbName).getConnection |
There was a problem hiding this comment.
Opening a raw connection via MockTexeraDB.getDBInstance.getDatabase(...) works (jOOQ qualifies the schema), but it couples the test to the trait's internals — a small newRawConnection() helper on the trait would read cleaner, and the same pattern repeats at 5 other sites in this file. Minor: getConnection + setAutoCommit(false) run before the try, so an early throw would leak the connection.
What changes were proposed in this PR?
Closes #6063 by completely refactoring the
MockTexeraDBinfrastructure to transition from a distributed process architecture to a centralized, shared-singleton architecture.MockTexeraDBtrait into a companionobject MockTexeraDBsingleton and a lightweighttrait MockTexeraDBinterface wrapper. The backgroundEmbeddedPostgresdatabase engine process is now booted exactly once per JVM lifecycle via synchronized initialization. Individual test suites dynamically provision an isolated schema sandbox via cheapCREATE DATABASEqueries instead of spawning heavy process-level instances.ddlScript), preventing parallel test suites from thrashing disk resources over the samesql/texera_ddl.sqlfile.PSQLException: Cannot commit when autoCommit is enabled). RefactoredSqlServer.withTransactionto safely intercept test configurations, bypassing rigid jOOQ transaction managers to manually cycle and restoreautoCommitstates cleanly on the active connection handle without leaking parameters across tests.Performance Impact
Moving database instantiation, parsing, and regex evaluation out of the distributed trait lifecycle and into a centralized static memory layer resulted in a dramatic reduction in test suite runtime.
Any related issues, documentation, discussions?
Follow up on #4525 where a refactor of
MockTexeraDBwas needed to allow for future concurrent execution of test suites. Structural baseline was cherry-picked and heavily adapted from #4527, specifically commitbd93161written by @Yicong-Huang.How was this PR tested?
Verified by running full, clean compilations and local stress tests on JDK 17 to validate that the structural split preserves total API backward compatibility across the codebase.
Additionally validated that high-concurrency race conditions (such as the asynchronous multi-threaded tests inside
DatasetResourceSpec) converge gracefully without experiencing deadlocks, thread crashes, or connection state leakage. All 1,600+ test suite assertions across the system are green.Run configuration command used to validate the project modules:
sbt clean compile testWas this PR authored or co-authored using generative AI tooling?
Yes. Generative AI (Gemini) was utilized as an external assistant to analyze complex multi-threaded JDBC/jOOQ connection trace errors and help structure the safe connection lifecycle fallback logic in SqlServer.scala. No direct IDE-integrated automation tools were used to auto-generate codebase files.